# -*- coding: utf-8 -*-
# @Author  : Zhang.Jingyi
import numpy as np
import json
import argparse
import os
import pickle

parser = argparse.ArgumentParser()
parser.add_argument('--json_file', type=str, required=True)
parser.add_argument('--out_file', type=str, required=True)
parser.add_argument('--face_file', type=str, required=True)
args = parser.parse_args()


def create_output(json_file, out_file, face_indx):
    with open(json_file, 'r') as load_f:
        load_dict = json.load(load_f)
    vertices = load_dict['pred_vertices_smpl']
    vertice = np.array(vertices)
    # content = np.vstack(vertice,face_indx)
    with open(out_file, "ab") as zjy_f:
        np.savetxt(zjy_f, vertice, fmt='%f %f %f')
        np.savetxt(zjy_f, face_indx, fmt='%d %d %d %d')
    ply_header = '''ply
format ascii 1.0
element vertex 6890
property float x
property float y
property float z
element face 13776
property list uchar int vertex_indices
end_header
    '''
    with open(out_file, 'r+') as f:
        old = f.read()
        f.seek(0)
        f.write(ply_header)
        f.write(old)


if __name__ == '__main__':
    with open(args.face_file, 'rb') as file:
        face_indx = pickle.load(file, encoding='iso-8859-1')['f']
    face_1 = np.ones((face_indx.shape[0], 1))
    face_1 *= 3
    # print(face_1.shape)
    # print(face_indx.shape)
    face = np.hstack((face_1, face_indx))
    create_output(args.json_file, args.out_file, face.astype(int))

